import pygame
import sys
import random
import math
from heapq import heappush, heappop

# 初始化 Pygame
pygame.init()

# 屏幕设置
WIDTH, HEIGHT = 1000, 600
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("算法可视化演示")

# 颜色定义
BACKGROUND = (10, 20, 30)
WHITE = (255, 255, 255)
RED = (255, 50, 50)
GREEN = (50, 255, 50)
BLUE = (50, 150, 255)
YELLOW = (255, 255, 0)
PURPLE = (180, 70, 240)
CYAN = (0, 200, 200)

# 字体
font = pygame.font.SysFont("microsoftyahei", 24)
title_font = pygame.font.SysFont("microsoftyahei", 36, bold=True)

class AlgorithmVisualizer:
    def __init__(self):
        self.mode = "sort"
        self.sort_algo = "bubble"
        self.path_algo = "dfs"
        self.array = []
        self.array_size = 50
        self.generate_array()
        
        # 网格设置
        self.grid_size = 15
        self.grid = [[0 for _ in range(self.grid_size)] for _ in range(self.grid_size)]
        self.start_pos = (2, 2)
        self.end_pos = (12, 12)
        self.grid[self.start_pos[1]][self.start_pos[0]] = 2  # 起点
        self.grid[self.end_pos[1]][self.end_pos[0]] = 3      # 终点
        
        # 分形设置
        self.fractal_depth = 4
        self.fractal_type = "tree"
                # 这些属性用于在排序可视化过程中高亮显示特定元素
        self.highlight_idx = -1    # 当前高亮索引
        self.compared_idx = -1     # 当前比较索引
        self.sorted_idx = -1       # 已排序部分的边界索引
    def generate_array(self):
        self.array = [random.randint(10, 100) for _ in range(self.array_size)]
        
    def draw_menu(self):
        # 绘制模式选择
        pygame.draw.rect(screen, (40, 50, 60), (10, 10, 980, 60), border_radius=5)
        
        modes = ["排序算法", "路径查找", "分形几何"]
        for i, mode in enumerate(modes):
            color = BLUE if self.mode == mode else WHITE
            text = font.render(mode, True, color)
            screen.blit(text, (100 + i * 300, 30))
            
        # 绘制算法选择
        pygame.draw.rect(screen, (40, 50, 60), (10, 80, 980, 50), border_radius=5)
        
        if self.mode == "排序算法":
            algorithms = ["冒泡排序", "选择排序", "插入排序", "快速排序"]
            for i, algo in enumerate(algorithms):
                color = GREEN if self.sort_algo == algo else WHITE
                text = font.render(algo, True, color)
                screen.blit(text, (100 + i * 200, 95))
                
        elif self.mode == "路径查找":
            algorithms = ["深度优先搜索", "广度优先搜索", "A*算法", " Dijkstra算法"]
            for i, algo in enumerate(algorithms):
                color = GREEN if self.path_algo == algo else WHITE
                text = font.render(algo, True, color)
                screen.blit(text, (100 + i * 200, 95))
                
        else:
            algorithms = ["分形树", "科赫雪花", "谢尔宾斯基三角形", "曼德勃罗集"]
            for i, algo in enumerate(algorithms):
                color = GREEN if self.fractal_type == algo else WHITE
                text = font.render(algo, True, color)
                screen.blit(text, (100 + i * 200, 95))
        
        # 绘制控制按钮
        pygame.draw.rect(screen, (40, 50, 60), (10, 140, 980, 50), border_radius=5)
        
        buttons = ["生成新数据", "开始可视化", "增加深度/大小", "减少深度/大小"]
        for i, button in enumerate(buttons):
            pygame.draw.rect(screen, BLUE, (100 + i * 200, 145, 180, 40), border_radius=5)
            text = font.render(button, True, WHITE)
            screen.blit(text, (110 + i * 200, 155))
    
    def draw_array(self):
        bar_width = (WIDTH - 40) // len(self.array)
        for i, val in enumerate(self.array):
            color = BLUE
            if i == self.highlight_idx:
                color = RED
            elif i == self.compared_idx:
                color = YELLOW
            elif i > self.sorted_idx:
                color = GREEN
                
            pygame.draw.rect(screen, color, (20 + i * bar_width, HEIGHT - 20 - val * 4, bar_width - 2, val * 4))
            
    def draw_grid(self):
        cell_size = min(500 // self.grid_size, 30)
        start_x = (WIDTH - cell_size * self.grid_size) // 2
        start_y = 200
        
        for y in range(self.grid_size):
            for x in range(self.grid_size):
                rect = pygame.Rect(start_x + x * cell_size, start_y + y * cell_size, cell_size, cell_size)
                
                if self.grid[y][x] == 1:  # 墙壁
                    pygame.draw.rect(screen, (50, 50, 70), rect)
                elif self.grid[y][x] == 2:  # 起点
                    pygame.draw.rect(screen, GREEN, rect)
                elif self.grid[y][x] == 3:  # 终点
                    pygame.draw.rect(screen, RED, rect)
                elif self.grid[y][x] == 4:  # 路径
                    pygame.draw.rect(screen, BLUE, rect)
                elif self.grid[y][x] == 5:  # 已访问
                    pygame.draw.rect(screen, (70, 70, 100), rect)
                else:  # 空单元格
                    pygame.draw.rect(screen, (30, 30, 40), rect)
                
                pygame.draw.rect(screen, (20, 20, 30), rect, 1)
    
    def draw_fractal_tree(self, x1, y1, angle, depth):
        if depth == 0:
            return
            
        branch_len = 10 * depth
        x2 = x1 + int(math.cos(math.radians(angle)) * branch_len)
        y2 = y1 - int(math.sin(math.radians(angle)) * branch_len)
        
        color = (50, 200 - depth * 30, 50)
        pygame.draw.line(screen, color, (x1, y1), (x2, y2), depth)
        
        self.draw_fractal_tree(x2, y2, angle - 20, depth - 1)
        self.draw_fractal_tree(x2, y2, angle + 20, depth - 1)
    
    def draw_koch_snowflake(self, x1, y1, x2, y2, depth):
        if depth == 0:
            pygame.draw.line(screen, CYAN, (x1, y1), (x2, y2), 2)
            return
            
        dx = x2 - x1
        dy = y2 - y1
        
        # 计算三等分点
        x3 = x1 + dx / 3
        y3 = y1 + dy / 3
        x4 = x1 + 2 * dx / 3
        y4 = y1 + 2 * dy / 3
        
        # 计算凸点
        x5 = x3 + (x4 - x3) * math.cos(math.radians(60)) - (y4 - y3) * math.sin(math.radians(60))
        y5 = y3 + (x4 - x3) * math.sin(math.radians(60)) + (y4 - y3) * math.cos(math.radians(60))
        
        # 递归绘制
        self.draw_koch_snowflake(x1, y1, x3, y3, depth - 1)
        self.draw_koch_snowflake(x3, y3, x5, y5, depth - 1)
        self.draw_koch_snowflake(x5, y5, x4, y4, depth - 1)
        self.draw_koch_snowflake(x4, y4, x2, y2, depth - 1)
    
    def draw_sierpinski(self, x1, y1, x2, y2, x3, y3, depth):
        if depth == 0:
            pygame.draw.polygon(screen, PURPLE, [(x1, y1), (x2, y2), (x3, y3)])
            return
            
        # 计算中点
        x12 = (x1 + x2) / 2
        y12 = (y1 + y2) / 2
        x23 = (x2 + x3) / 2
        y23 = (y2 + y3) / 2
        x31 = (x3 + x1) / 2
        y31 = (y3 + y1) / 2
        
        # 递归绘制
        self.draw_sierpinski(x1, y1, x12, y12, x31, y31, depth - 1)
        self.draw_sierpinski(x12, y12, x2, y2, x23, y23, depth - 1)
        self.draw_sierpinski(x31, y31, x23, y23, x3, y3, depth - 1)
    
    def draw_mandelbrot(self):
        for x in range(300):
            for y in range(300):
                cx = (x - 150) / 50
                cy = (y - 150) / 50
                
                zx, zy = 0, 0
                for i in range(20):
                    zx, zy = zx*zx - zy*zy + cx, 2*zx*zy + cy
                    if zx*zx + zy*zy > 4:
                        break
                
                # color = (i * 12, i * 8, i * 20)
                # pygame.draw.rect(screen, color, (350 + x, 200 + y, 1, 1))
                    # 根据迭代次数生成颜色值
                # 修复：确保颜色值在0-255范围内，避免无效颜色参数错误
                color = (min(i * 12, 255), min(i * 8, 255), min(i * 20, 255))
                # 在屏幕上绘制该像素点
                pygame.draw.rect(screen, color, (350 + x, 200 + y, 1, 1))

    def draw_fractal(self):
        if self.fractal_type == "分形树":
            self.draw_fractal_tree(WIDTH // 2, HEIGHT - 50, 90, self.fractal_depth)
        elif self.fractal_type == "科赫雪花":
            size = 150
            x1, y1 = WIDTH // 2 - size, HEIGHT // 2 + size // 2
            x2, y2 = WIDTH // 2 + size, HEIGHT // 2 + size // 2
            x3, y3 = WIDTH // 2, HEIGHT // 2 - size // 2
            
            self.draw_koch_snowflake(x1, y1, x2, y2, self.fractal_depth)
            self.draw_koch_snowflake(x2, y2, x3, y3, self.fractal_depth)
            self.draw_koch_snowflake(x3, y3, x1, y1, self.fractal_depth)
        elif self.fractal_type == "谢尔宾斯基三角形":
            size = 200
            x1, y1 = WIDTH // 2, HEIGHT // 2 - size // 2
            x2, y2 = WIDTH // 2 - size // 2, HEIGHT // 2 + size // 2
            x3, y3 = WIDTH // 2 + size // 2, HEIGHT // 2 + size // 2
            
            self.draw_sierpinski(x1, y1, x2, y2, x3, y3, self.fractal_depth)
        else:
            self.draw_mandelbrot()
    
    def draw(self):
        screen.fill(BACKGROUND)
        self.draw_menu()
        
        title = title_font.render(f"{self.mode} - {self.sort_algo if self.mode == '排序算法' else self.path_algo if self.mode == '路径查找' else self.fractal_type}", True, YELLOW)
        screen.blit(title, (WIDTH // 2 - title.get_width() // 2, 160))
        
        if self.mode == "排序算法":
            self.draw_array()
        elif self.mode == "路径查找":
            self.draw_grid()
        else:
            self.draw_fractal()
    
    def bubble_sort_step(self):
        n = len(self.array)
        for i in range(n - 1 - self.sorted_idx):
            self.highlight_idx = i
            self.compared_idx = i + 1
            
            if self.array[i] > self.array[i + 1]:
                self.array[i], self.array[i + 1] = self.array[i + 1], self.array[i]
                
            return False
        
        self.sorted_idx += 1
        if self.sorted_idx >= n - 1:
            self.sorted_idx = n
            return True
        return False
    
    def selection_sort_step(self):
        n = len(self.array)
        if self.sorted_idx < n - 1:
            min_idx = self.sorted_idx
            for i in range(self.sorted_idx + 1, n):
                self.highlight_idx = i
                self.compared_idx = min_idx
                
                if self.array[i] < self.array[min_idx]:
                    min_idx = i
                    
            self.array[self.sorted_idx], self.array[min_idx] = self.array[min_idx], self.array[self.sorted_idx]
            self.sorted_idx += 1
            return False
        return True
    
    def insertion_sort_step(self):
        n = len(self.array)
        if self.sorted_idx < n:
            key = self.array[self.sorted_idx]
            j = self.sorted_idx - 1
            
            self.highlight_idx = self.sorted_idx
            self.compared_idx = j
            
            while j >= 0 and self.array[j] > key:
                self.array[j + 1] = self.array[j]
                j -= 1
                
            self.array[j + 1] = key
            self.sorted_idx += 1
            
            if self.sorted_idx >= n:
                return True
        return False
    
    def quick_sort_step(self):
        # 简化的快速排序步骤
        if not hasattr(self, 'quick_stack'):
            self.quick_stack = [(0, len(self.array) - 1)]
            self.sorted_idx = -1
            
        if self.quick_stack:
            low, high = self.quick_stack.pop()
            if low < high:
                pivot = self.array[high]
                i = low - 1
                
                for j in range(low, high):
                    self.highlight_idx = j
                    self.compared_idx = high
                    
                    if self.array[j] <= pivot:
                        i += 1
                        self.array[i], self.array[j] = self.array[j], self.array[i]
                        
                self.array[i + 1], self.array[high] = self.array[high], self.array[i + 1]
                pi = i + 1
                
                self.quick_stack.append((low, pi - 1))
                self.quick_stack.append((pi + 1, high))
                
            return False
        return True
    
    def dfs_step(self):
        if not hasattr(self, 'dfs_stack'):
            self.dfs_stack = [self.start_pos]
            self.dfs_visited = set([self.start_pos])
            self.grid[self.start_pos[1]][self.start_pos[0]] = 5  # 标记为已访问
            
        if self.dfs_stack:
            x, y = self.dfs_stack.pop()
            
            if (x, y) == self.end_pos:
                return True
                
            for dx, dy in [(0, 1), (1, 0), (0, -1), (-1, 0)]:
                nx, ny = x + dx, y + dy
                if 0 <= nx < self.grid_size and 0 <= ny < self.grid_size:
                    if self.grid[ny][nx] != 1 and (nx, ny) not in self.dfs_visited:
                        self.dfs_stack.append((nx, ny))
                        self.dfs_visited.add((nx, ny))
                        if (nx, ny) != self.end_pos:
                            self.grid[ny][nx] = 5  # 标记为已访问
                            
            return False
        return True
    
    def bfs_step(self):
        if not hasattr(self, 'bfs_queue'):
            self.bfs_queue = [self.start_pos]
            self.bfs_visited = set([self.start_pos])
            self.grid[self.start_pos[1]][self.start_pos[0]] = 5  # 标记为已访问
            
        if self.bfs_queue:
            x, y = self.bfs_queue.pop(0)
            
            if (x, y) == self.end_pos:
                return True
                
            for dx, dy in [(0, 1), (1, 0), (0, -1), (-1, 0)]:
                nx, ny = x + dx, y + dy
                if 0 <= nx < self.grid_size and 0 <= ny < self.grid_size:
                    if self.grid[ny][nx] != 1 and (nx, ny) not in self.bfs_visited:
                        self.bfs_queue.append((nx, ny))
                        self.bfs_visited.add((nx, ny))
                        if (nx, ny) != self.end_pos:
                            self.grid[ny][nx] = 5  # 标记为已访问
                            
            return False
        return True
    
    def handle_click(self, pos):
        x, y = pos
        
        # 检查模式选择
        if 20 < y < 80:
            if 100 < x < 250:
                self.mode = "排序算法"
            elif 400 < x < 550:
                self.mode = "路径查找"
            elif 700 < x < 850:
                self.mode = "分形几何"
            return
        
        # 检查算法选择
        if 80 < y < 130:
            if self.mode == "排序算法":
                if 100 < x < 250:
                    self.sort_algo = "冒泡排序"
                elif 300 < x < 450:
                    self.sort_algo = "选择排序"
                elif 500 < x < 650:
                    self.sort_algo = "插入排序"
                elif 700 < x < 850:
                    self.sort_algo = "快速排序"
            elif self.mode == "路径查找":
                if 100 < x < 250:
                    self.path_algo = "深度优先搜索"
                elif 300 < x < 450:
                    self.path_algo = "广度优先搜索"
                elif 500 < x < 650:
                    self.path_algo = "A*算法"
                elif 700 < x < 850:
                    self.path_algo = " Dijkstra算法"
            else:
                if 100 < x < 250:
                    self.fractal_type = "分形树"
                elif 300 < x < 450:
                    self.fractal_type = "科赫雪花"
                elif 500 < x < 650:
                    self.fractal_type = "谢尔宾斯基三角形"
                elif 700 < x < 850:
                    self.fractal_type = "曼德勃罗集"
            return
        
        # 检查控制按钮
        if 140 < y < 190:
            if 100 < x < 280:
                if self.mode == "排序算法":
                    self.generate_array()
                elif self.mode == "路径查找":
                    # 重置网格
                    self.grid = [[0 for _ in range(self.grid_size)] for _ in range(self.grid_size)]
                    # 随机添加墙壁
                    for i in range(self.grid_size):
                        for j in range(self.grid_size):
                            if random.random() < 0.3 and (i, j) != self.start_pos and (i, j) != self.end_pos:
                                self.grid[j][i] = 1
                    self.grid[self.start_pos[1]][self.start_pos[0]] = 2
                    self.grid[self.end_pos[1]][self.end_pos[0]] = 3
            elif 300 < x < 480:
                self.visualizing = True
                self.sorted_idx = 0
                self.highlight_idx = -1
                self.compared_idx = -1
                
                # 清除排序算法属性
                if hasattr(self, 'quick_stack'):
                    delattr(self, 'quick_stack')
            elif 500 < x < 680:
                if self.mode == "排序算法":
                    self.array_size = min(100, self.array_size + 10)
                    self.generate_array()
                elif self.mode == "路径查找":
                    self.grid_size = min(30, self.grid_size + 2)
                    self.grid = [[0 for _ in range(self.grid_size)] for _ in range(self.grid_size)]
                    self.start_pos = (2, 2)
                    self.end_pos = (self.grid_size - 3, self.grid_size - 3)
                    self.grid[self.start_pos[1]][self.start_pos[0]] = 2
                    self.grid[self.end_pos[1]][self.end_pos[0]] = 3
                else:
                    self.fractal_depth = min(8, self.fractal_depth + 1)
            elif 700 < x < 880:
                if self.mode == "排序算法":
                    self.array_size = max(10, self.array_size - 10)
                    self.generate_array()
                elif self.mode == "路径查找":
                    self.grid_size = max(10, self.grid_size - 2)
                    self.grid = [[0 for _ in range(self.grid_size)] for _ in range(self.grid_size)]
                    self.start_pos = (2, 2)
                    self.end_pos = (self.grid_size - 3, self.grid_size - 3)
                    self.grid[self.start_pos[1]][self.start_pos[0]] = 2
                    self.grid[self.end_pos[1]][self.end_pos[0]] = 3
                else:
                    self.fractal_depth = max(1, self.fractal_depth - 1)
        
        # 网格点击处理
        if self.mode == "路径查找" and y > 200:
            cell_size = min(500 // self.grid_size, 30)
            start_x = (WIDTH - cell_size * self.grid_size) // 2
            start_y = 200
            
            grid_x = (x - start_x) // cell_size
            grid_y = (y - start_y) // cell_size
            
            if 0 <= grid_x < self.grid_size and 0 <= grid_y < self.grid_size:
                if (grid_x, grid_y) != self.start_pos and (grid_x, grid_y) != self.end_pos:
                    if self.grid[grid_y][grid_x] == 0:  # 空单元格
                        self.grid[grid_y][grid_x] = 1  # 设置为墙壁
                    elif self.grid[grid_y][grid_x] == 1:  # 墙壁
                        self.grid[grid_y][grid_x] = 0  # 设置为空单元格
    def update(self):
        # 更新可视化状态，处理算法的每一步执行
        if hasattr(self, 'visualizing') and self.visualizing:
            if self.mode == "排序算法":
                # 根据当前选择的排序算法执行相应的步骤
                if self.sort_algo == "冒泡排序":
                    done = self.bubble_sort_step()
                elif self.sort_algo == "选择排序":
                    done = self.selection_sort_step()
                elif self.sort_algo == "插入排序":
                    done = self.insertion_sort_step()
                else:  # 快速排序
                    done = self.quick_sort_step()
                
                # 如果算法执行完成，停止可视化
                if done:
                    self.visualizing = False
                    
            elif self.mode == "路径查找":
                # 根据当前选择的路径查找算法执行相应的步骤
                if self.path_algo == "深度优先搜索":
                    done = self.dfs_step()
                elif self.path_algo == "广度优先搜索":
                    done = self.bfs_step()
                # 其他路径查找算法实现类似，为简化代码此处省略
                else:
                    # 对于未实现的算法，直接标记为完成
                    done = True
                
                # 如果算法执行完成，停止可视化
                if done:
                    self.visualizing = False

# 创建可视化器实例
visualizer = AlgorithmVisualizer()

# 主循环
clock = pygame.time.Clock()
running = True

while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        elif event.type == pygame.MOUSEBUTTONDOWN:
            visualizer.handle_click(event.pos)
    
    visualizer.update()
    visualizer.draw()
    
    pygame.display.flip()
    clock.tick(30)

pygame.quit()
sys.exit()